package top.hjcwzx.wrench.limiter.strategy.algorithm;

import top.hjcwzx.wrench.limiter.config.RateLimiterLeakyBucketAutoProperties;
import top.hjcwzx.wrench.limiter.config.RateLimiterMemoryAutoProperties;
import top.hjcwzx.wrench.limiter.config.RateLimiterTokenBucketAutoProperties;
import top.hjcwzx.wrench.limiter.enums.AlgorithmStrategyEnum;
import top.hjcwzx.wrench.limiter.memory.ExpirableKeyStore;
import top.hjcwzx.wrench.limiter.strategy.algorithm.impl.LeakyBucketRateLimiterStrategy;
import top.hjcwzx.wrench.limiter.strategy.algorithm.impl.TokenBucketRateLimiterStrategy;

import javax.annotation.PostConstruct;
import javax.annotation.Resource;
import java.util.HashMap;
import java.util.Map;

//限流策略上下文
public class RateLimiterAlgorithmStrategyContext {

    private final Map<AlgorithmStrategyEnum, RateLimiterAlgorithmStrategy> algorithmMap = new HashMap<>();

    @Resource
    private RateLimiterTokenBucketAutoProperties tokenBucketAutoProperties;

    @Resource
    private RateLimiterMemoryAutoProperties memoryAutoProperties;

    @Resource
    private RateLimiterLeakyBucketAutoProperties leakyBucketAutoProperties;


    @PostConstruct
    public void init() {
        ExpirableKeyStore.EvictionPolicy evictionPolicy = memoryAutoProperties.getEvictionPolicy();
        Long expireAfterAccessMs = memoryAutoProperties.getExpireAfterAccessMs();
        Integer maxCacheSize = memoryAutoProperties.getMaxCacheSize();

        Integer bucketTokenCapacity = tokenBucketAutoProperties.getBucketTokenCapacity();
        Double refillRatePerSecond = tokenBucketAutoProperties.getRefillRatePerSecond();

        Integer bucketLeakyCapacity = leakyBucketAutoProperties.getBucketLeakyCapacity();
        Double leakRate = leakyBucketAutoProperties.getLeakRate();

        //令牌桶算法
        algorithmMap.put(
                AlgorithmStrategyEnum.TOKEN_BUCKET,
                new TokenBucketRateLimiterStrategy(
                        //内存淘汰策略
                        evictionPolicy,
                        //过期时间
                        expireAfterAccessMs,
                        //最大key数量
                        maxCacheSize,
                        //桶中令牌的容量
                        bucketTokenCapacity,
                        //每秒补充的令牌数
                        refillRatePerSecond
                )
        );
        // 漏斗算法
        algorithmMap.put(
                AlgorithmStrategyEnum.LEAKY_BUCKET,
                new LeakyBucketRateLimiterStrategy(
                        //内存淘汰策略
                        evictionPolicy,
                        //过期时间
                        expireAfterAccessMs,
                        //最大key数量
                        maxCacheSize,
                        //漏斗的容量
                        bucketLeakyCapacity,
                        //每秒漏斗速率
                        leakRate
                )
        );

    }

    public RateLimiterAlgorithmStrategy getAlgorithm(AlgorithmStrategyEnum name) {
        RateLimiterAlgorithmStrategy strategy = algorithmMap.get(name);
        return null == strategy ? algorithmMap.get(AlgorithmStrategyEnum.TOKEN_BUCKET) : strategy;
    }
}
